import ctypes
from ctypes import wintypes
import os
import time
import threading
import tkinter as tk
from tkinter import ttk, messagebox, scrolledtext


class CH341I2CGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("通用CH341 I2C助手")
        self.root.geometry("900x750")
        self.root.resizable(True, True)

        # 全局变量
        self.tool = CH341I2CTool()
        self.loop_thread = None
        self.loop_running = False
        self.current_loop = 0

        # 创建界面
        self._create_widgets()
        self._update_address_display()
        self._toggle_reg_addr_input()

    def _create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 1. 设备控制区
        control_frame = ttk.LabelFrame(main_frame, text="设备控制", padding="10")
        control_frame.pack(fill=tk.X, pady=(0, 10))

        self.init_btn = ttk.Button(control_frame, text="初始化设备", command=self._init_device)
        self.init_btn.grid(row=0, column=0, padx=5)
        self.close_btn = ttk.Button(control_frame, text="关闭设备", command=self._close_device, state=tk.DISABLED)
        self.close_btn.grid(row=0, column=1, padx=5)
        
        # I2C速率选择
        ttk.Label(control_frame, text="I2C速率：").grid(row=0, column=2, padx=5)
        self.speed_var = tk.StringVar(value="100KHz")
        speed_combobox = ttk.Combobox(
            control_frame, 
            textvariable=self.speed_var,
            values=["10KHz", "50KHz", "100KHz", "400KHz"],
            state="readonly",
            width=8
        )
        speed_combobox.grid(row=0, column=3, padx=5)
        
        self.status_var = tk.StringVar(value="状态：未初始化")
        ttk.Label(control_frame, textvariable=self.status_var).grid(row=0, column=4, padx=20, sticky=tk.W)

        # 2. 参数输入区
        param_frame = ttk.LabelFrame(main_frame, text="参数设置", padding="10")
        param_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=False, pady=(0, 10), ipadx=5)

        # 设备地址设置
        self.use_dev_addr = tk.BooleanVar(value=True)
        ttk.Checkbutton(
            param_frame, 
            text="使用设备地址", 
            variable=self.use_dev_addr,
            command=self._update_address_display
        ).grid(row=0, column=0, columnspan=2, sticky=tk.W, pady=5)
        
        ttk.Label(param_frame, text="7位设备地址 (0x)：").grid(row=1, column=0, sticky=tk.W)
        self.dev_addr_var = tk.StringVar(value="48")
        self.dev_addr_entry = ttk.Entry(param_frame, textvariable=self.dev_addr_var, width=10)
        self.dev_addr_entry.grid(row=1, column=1, pady=2)
        
        # 8位地址显示
        ttk.Label(param_frame, text="8位写地址：").grid(row=2, column=0, sticky=tk.W)
        self.dev_addr_write_8bit = tk.StringVar(value="0x90")
        ttk.Label(param_frame, textvariable=self.dev_addr_write_8bit).grid(row=2, column=1, sticky=tk.W)
        
        ttk.Label(param_frame, text="8位读地址：").grid(row=3, column=0, sticky=tk.W)
        self.dev_addr_read_8bit = tk.StringVar(value="0x91")
        ttk.Label(param_frame, textvariable=self.dev_addr_read_8bit).grid(row=3, column=1, sticky=tk.W)

        # 寄存器地址设置
        self.use_reg_addr = tk.BooleanVar(value=True)
        ttk.Checkbutton(
            param_frame, 
            text="使用寄存器地址", 
            variable=self.use_reg_addr,
            command=self._toggle_reg_addr_input
        ).grid(row=4, column=0, columnspan=2, sticky=tk.W, pady=5)
        
        ttk.Label(param_frame, text="寄存器地址 (0x)：").grid(row=5, column=0, sticky=tk.W)
        self.reg_addr_var = tk.StringVar(value="00")
        self.reg_addr_entry = ttk.Entry(param_frame, textvariable=self.reg_addr_var, width=10)
        self.reg_addr_entry.grid(row=5, column=1, pady=2)

        # 数据输入区
        ttk.Label(param_frame, text="数据 (空格分隔)：").grid(row=6, column=0, sticky=tk.NW, pady=5)
        self.data_var = tk.Text(param_frame, width=20, height=5)
        self.data_var.grid(row=6, column=1, pady=5)
        self.data_var.insert(tk.END, "AA BB")

        # 3. 操作区
        op_frame = ttk.LabelFrame(main_frame, text="操作", padding="10")
        op_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=False, pady=(0, 10), ipadx=5)

        # 通用读写按钮
        self.read_btn = ttk.Button(op_frame, text="读取数据", command=self._read_data)
        self.read_btn.grid(row=0, column=0, padx=5, pady=5, sticky=tk.EW)
        self.write_btn = ttk.Button(op_frame, text="写入数据", command=self._write_data)
        self.write_btn.grid(row=0, column=1, padx=5, pady=5, sticky=tk.EW)

        # 读取长度设置
        ttk.Label(op_frame, text="读取长度：").grid(row=1, column=0, sticky=tk.W)
        self.read_length_var = tk.StringVar(value="2")
        ttk.Entry(op_frame, textvariable=self.read_length_var, width=8).grid(row=1, column=1, sticky=tk.W)

        # 循环发送设置
        ttk.Label(op_frame, text="循环发送设置：").grid(row=2, column=0, columnspan=2, sticky=tk.W, pady=10)
        ttk.Label(op_frame, text="次数 (0=无限)：").grid(row=3, column=0, sticky=tk.W)
        self.loop_count_var = tk.StringVar(value="10")
        ttk.Entry(op_frame, textvariable=self.loop_count_var, width=8).grid(row=3, column=1, sticky=tk.W)
        ttk.Label(op_frame, text="间隔 (秒)：").grid(row=4, column=0, sticky=tk.W)
        self.loop_interval_var = tk.StringVar(value="0.5")
        ttk.Entry(op_frame, textvariable=self.loop_interval_var, width=8).grid(row=4, column=1, sticky=tk.W)

        # 循环控制按钮
        self.start_loop_btn = ttk.Button(op_frame, text="开始循环", command=self._start_loop, state=tk.DISABLED)
        self.start_loop_btn.grid(row=5, column=0, padx=5, pady=10, sticky=tk.EW)
        self.stop_loop_btn = ttk.Button(op_frame, text="停止循环", command=self._stop_loop, state=tk.DISABLED)
        self.stop_loop_btn.grid(row=5, column=1, padx=5, pady=10, sticky=tk.EW)

        # 4. 日志区
        log_frame = ttk.LabelFrame(main_frame, text="操作日志", padding="10")
        log_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 10))
        
        log_control_frame = ttk.Frame(log_frame)
        log_control_frame.pack(fill=tk.X, pady=(0, 5))
        
        self.clear_log_btn = ttk.Button(log_control_frame, text="清空日志", command=self._clear_log)
        self.clear_log_btn.pack(side=tk.RIGHT)

        self.log_text = scrolledtext.ScrolledText(log_frame, wrap=tk.WORD, height=15)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        self.log_text.config(state=tk.DISABLED)

        # 5. 状态栏
        self.status_bar = ttk.Label(self.root, text="就绪", relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)

    def _update_address_display(self):
        if self.use_dev_addr.get():
            self.dev_addr_entry.config(state=tk.NORMAL)
            try:
                addr_7bit = int(self.dev_addr_var.get(), 16)
                if 0 <= addr_7bit <= 0x7F:
                    addr_write_8bit = (addr_7bit << 1) & 0xFE
                    addr_read_8bit = (addr_7bit << 1) | 0x01
                    self.dev_addr_write_8bit.set(f"0x{addr_write_8bit:02X}")
                    self.dev_addr_read_8bit.set(f"0x{addr_read_8bit:02X}")
            except:
                pass
        else:
            self.dev_addr_entry.config(state=tk.DISABLED)
            self.dev_addr_write_8bit.set("0x00 (广播)")
            self.dev_addr_read_8bit.set("0x01 (广播读)")

    def _toggle_reg_addr_input(self):
        if self.use_reg_addr.get():
            self.reg_addr_entry.config(state=tk.NORMAL)
        else:
            self.reg_addr_entry.config(state=tk.DISABLED)

    def _clear_log(self):
        self.log_text.config(state=tk.NORMAL)
        self.log_text.delete(1.0, tk.END)
        self.log_text.config(state=tk.DISABLED)
        self._log("日志已清空")

    def _get_speed_code(self):
        speed_map = {
            "10KHz": 0x00,
            "50KHz": 0x02,
            "100KHz": 0x01,
            "400KHz": 0x03
        }
        return speed_map.get(self.speed_var.get(), 0x01)

    def _init_device(self):
        speed_code = self._get_speed_code()
        if self.tool.init(speed_code):
            self.status_var.set(f"状态：设备已连接（速率：{self.speed_var.get()}）")
            self.init_btn.config(state=tk.DISABLED)
            self.close_btn.config(state=tk.NORMAL)
            self.start_loop_btn.config(state=tk.NORMAL)
            self._log(f"设备初始化成功，速率：{self.speed_var.get()}")
        else:
            messagebox.showerror("错误", "设备初始化失败，请检查连接和驱动！")
            self._log("设备初始化失败")

    def _close_device(self):
        self._stop_loop()
        self.tool.close()
        self.status_var.set("状态：设备已关闭")
        self.init_btn.config(state=tk.NORMAL)
        self.close_btn.config(state=tk.DISABLED)
        self.start_loop_btn.config(state=tk.DISABLED)
        self._log("设备已关闭")

    def _read_data(self):
        try:
            dev_addr = int(self.dev_addr_var.get(), 16) if self.use_dev_addr.get() else 0x00
            use_reg = self.use_reg_addr.get()
            reg_addr = int(self.reg_addr_var.get(), 16) if use_reg else None
            read_length = int(self.read_length_var.get())
            if read_length < 1 or read_length > 255:
                raise ValueError("读取长度必须为1-255")

            data_list = self.tool.read_data(
                dev_addr=dev_addr,
                reg_addr=reg_addr,
                length=read_length,
                use_reg=use_reg
            )

            addr_info = f"设备0x{dev_addr:02X}" if self.use_dev_addr.get() else "广播地址"
            reg_info = f"，寄存器0x{reg_addr:02X}" if use_reg else ""
            if data_list is not None:
                data_str = " ".join([f"{d:02X}" for d in data_list])
                self._log(f"读取成功：{addr_info}{reg_info}，长度{read_length} → {data_str}")
                self.data_var.delete(1.0, tk.END)
                self.data_var.insert(tk.END, data_str)
            else:
                self._log(f"读取失败：{addr_info}{reg_info}，长度{read_length}")
        except Exception as e:
            self._log(f"读取错误：{str(e)}")
            messagebox.showerror("错误", f"参数错误：{str(e)}")

    def _write_data(self):
        try:
            dev_addr = int(self.dev_addr_var.get(), 16) if self.use_dev_addr.get() else 0x00
            use_reg = self.use_reg_addr.get()
            reg_addr = int(self.reg_addr_var.get(), 16) if use_reg else None
            data_str = self.data_var.get(1.0, tk.END).strip()
            data_list = [int(d, 16) for d in data_str.split()]
            if not data_list or len(data_list) > 255:
                raise ValueError("数据长度必须为1-255字节")

            success = self.tool.write_data(
                dev_addr=dev_addr,
                reg_addr=reg_addr,
                data_list=data_list,
                use_reg=use_reg
            )

            addr_info = f"设备0x{dev_addr:02X}" if self.use_dev_addr.get() else "广播地址"
            reg_info = f"，寄存器0x{reg_addr:02X}" if use_reg else ""
            self._log(f"写入{'成功' if success else '失败'}：{addr_info}{reg_info} → {data_str}（{len(data_list)}字节）")
        except Exception as e:
            self._log(f"写入错误：{str(e)}")
            messagebox.showerror("错误", f"参数错误：{str(e)}")

    def _start_loop(self):
        if self.loop_running:
            return
        try:
            dev_addr = int(self.dev_addr_var.get(), 16) if self.use_dev_addr.get() else 0x00
            use_reg = self.use_reg_addr.get()
            reg_addr = int(self.reg_addr_var.get(), 16) if use_reg else None
            data_str = self.data_var.get(1.0, tk.END).strip()
            data_list = [int(d, 16) for d in data_str.split()]
            if not data_list or len(data_list) > 255:
                raise ValueError("数据长度必须为1-255字节")
            loop_count = int(self.loop_count_var.get())
            interval = float(self.loop_interval_var.get())
            if interval < 0.01:
                interval = 0.01

            self.loop_running = True
            self.current_loop = 0
            self.start_loop_btn.config(state=tk.DISABLED)
            self.stop_loop_btn.config(state=tk.NORMAL)
            self._log(f"开始循环发送（次数：{loop_count if loop_count > 0 else '无限'}，间隔：{interval}秒）")

            self.loop_thread = threading.Thread(
                target=self._loop_send,
                args=(dev_addr, reg_addr, data_list, use_reg, loop_count, interval),
                daemon=True
            )
            self.loop_thread.start()
        except Exception as e:
            self._log(f"循环设置错误：{str(e)}")
            messagebox.showerror("错误", f"参数错误：{str(e)}")

    def _loop_send(self, dev_addr, reg_addr, data_list, use_reg, loop_count, interval):
        while self.loop_running:
            if loop_count > 0 and self.current_loop >= loop_count:
                self.root.after(0, self._stop_loop)
                break

            success = self.tool.write_data(
                dev_addr=dev_addr,
                reg_addr=reg_addr,
                data_list=data_list,
                use_reg=use_reg
            )
            
            self.current_loop += 1
            data_str = " ".join([f"0x{d:02X}" for d in data_list])
            status = "成功" if success else "失败"
            self.root.after(0, lambda: self._log(
                f"循环发送[{self.current_loop}] {status}：设备0x{dev_addr:02X}{'，寄存器0x'+f'{reg_addr:02X}' if use_reg else ''} → {data_str}"
            ))

            time.sleep(interval)

    def _stop_loop(self):
        self.loop_running = False
        if self.loop_thread is not None:
            self.loop_thread.join()
        self.start_loop_btn.config(state=tk.NORMAL)
        self.stop_loop_btn.config(state=tk.DISABLED)
        self._log(f"循环发送停止（总次数：{self.current_loop}）")

    def _log(self, msg):
        self.log_text.config(state=tk.NORMAL)
        self.log_text.insert(tk.END, f"[{time.strftime('%H:%M:%S')}] {msg}\n")
        self.log_text.see(tk.END)
        self.log_text.config(state=tk.DISABLED)


class CH341I2CTool:
    def __init__(self):
        self.initialized = False
        self.device_index = 0  # 默认设备索引
        try:
            script_dir = os.path.dirname(os.path.abspath(__file__))
            self.dll_path = os.path.join(script_dir, "CH341DLLA64.DLL")
            if not os.path.exists(self.dll_path):
                raise FileNotFoundError(f"未找到DLL文件: {self.dll_path}")
            self.ch341 = ctypes.WinDLL(self.dll_path)
            self._init_func_types()
        except Exception as e:
            self.ch341 = None
            print(f"DLL初始化失败: {e}")

    def _init_func_types(self):
        """初始化DLL函数参数和返回值类型"""
        # 打开设备
        self.ch341.CH341OpenDevice.argtypes = [ctypes.c_ulong]
        self.ch341.CH341OpenDevice.restype = wintypes.HANDLE

        # 关闭设备
        self.ch341.CH341CloseDevice.argtypes = [ctypes.c_ulong]
        self.ch341.CH341CloseDevice.restype = None

        # 设置I2C速率
        self.ch341.CH341SetStream.argtypes = [ctypes.c_ulong, ctypes.c_ulong]
        self.ch341.CH341SetStream.restype = ctypes.c_bool

        # 单字节读写（带寄存器）
        self.ch341.CH341WriteI2C.argtypes = [
            ctypes.c_ulong,  # 设备索引
            ctypes.c_ubyte,  # 7位设备地址
            ctypes.c_ubyte,  # 寄存器地址
            ctypes.c_ubyte   # 数据
        ]
        self.ch341.CH341WriteI2C.restype = ctypes.c_bool

        self.ch341.CH341ReadI2C.argtypes = [
            ctypes.c_ulong,  # 设备索引
            ctypes.c_ubyte,  # 7位设备地址
            ctypes.c_ubyte,  # 寄存器地址
            ctypes.POINTER(ctypes.c_ubyte)  # 接收数据指针
        ]
        self.ch341.CH341ReadI2C.restype = ctypes.c_bool

        # 流模式读写（无寄存器，用于纯数据传输）
        self.ch341.CH341StreamI2C.argtypes = [
            ctypes.c_ulong,  # 设备索引
            ctypes.c_ulong,  # 输出数据长度
            ctypes.POINTER(ctypes.c_ubyte),  # 输出数据
            ctypes.c_ulong,  # 输入数据长度
            ctypes.POINTER(ctypes.c_ubyte)   # 输入数据
        ]
        self.ch341.CH341StreamI2C.restype = ctypes.c_bool

    def init(self, i2c_speed=0x01):
        """初始化设备并设置速率"""
        if not self.ch341:
            return False
        # 打开设备
        handle = self.ch341.CH341OpenDevice(self.device_index)
        if handle == -1:  # INVALID_HANDLE_VALUE
            return False
        # 设置速率
        if not bool(self.ch341.CH341SetStream(self.device_index, i2c_speed)):
            self.ch341.CH341CloseDevice(self.device_index)
            return False
        self.initialized = True
        return True

    def write_data(self, dev_addr, data_list, reg_addr=None, use_reg=True):
        """
        写入数据：
        - 带寄存器：使用CH341WriteI2C（自动处理起始/停止位）
        - 无寄存器：使用CH341StreamI2C（不添加控制码）
        """
        if not self.initialized or not data_list:
            return False

        try:
            if use_reg and reg_addr is not None:
                # 带寄存器：单字节用WriteI2C，多字节用StreamI2C（不含控制码）
                if len(data_list) == 1:
                    # 单字节写入（自动处理时序）
                    result = self.ch341.CH341WriteI2C(
                        self.device_index,
                        ctypes.c_ubyte(dev_addr),
                        ctypes.c_ubyte(reg_addr),
                        ctypes.c_ubyte(data_list[0])
                    )
                    return bool(result)  # 转换为Python布尔值
                else:
                    # 多字节写入：设备地址(写) + 寄存器 + 数据
                    data_frame = [
                        (dev_addr << 1) & 0xFE,  # 8位写地址
                        reg_addr
                    ] + data_list
                    data_buf = (ctypes.c_ubyte * len(data_frame))(*data_frame)
                    result = self.ch341.CH341StreamI2C(
                        self.device_index,
                        len(data_frame),
                        data_buf,
                        0,
                        None
                    )
                    return bool(result)  # 转换为Python布尔值
            else:
                # 无寄存器：直接发送数据帧（设备地址 + 数据）
                data_frame = [
                    (dev_addr << 1) & 0xFE  # 8位写地址
                ] + data_list
                data_buf = (ctypes.c_ubyte * len(data_frame))(*data_frame)
                result = self.ch341.CH341StreamI2C(
                    self.device_index,
                    len(data_frame),
                    data_buf,
                    0,
                    None
                )
                return bool(result)  # 转换为Python布尔值
        except Exception as e:
            print(f"写入数据时发生错误: {e}")
            return False

    def read_data(self, dev_addr, length, reg_addr=None, use_reg=True):
        """
        读取数据：
        - 带寄存器：先写寄存器地址，再读数据
        - 无寄存器：直接读数据
        """
        if not self.initialized or length <= 0:
            return None

        try:
            if use_reg and reg_addr is not None:
                # 带寄存器：单字节用ReadI2C，多字节用StreamI2C
                if length == 1:
                    # 单字节读取（自动处理时序）
                    data = ctypes.c_ubyte(0)
                    result = self.ch341.CH341ReadI2C(
                        self.device_index,
                        ctypes.c_ubyte(dev_addr),
                        ctypes.c_ubyte(reg_addr),
                        ctypes.byref(data)
                    )
                    return [data.value] if bool(result) else None
                else:
                    # 多字节读取步骤1：写入寄存器地址
                    write_frame = [
                        (dev_addr << 1) & 0xFE,  # 8位写地址
                        reg_addr
                    ]
                    write_buf = (ctypes.c_ubyte * len(write_frame))(*write_frame)
                    result1 = self.ch341.CH341StreamI2C(
                        self.device_index,
                        len(write_frame),
                        write_buf,
                        0,
                        None
                    )
                    if not bool(result1):
                        return None
                    
                    # 步骤2：读取数据
                    read_frame = [(dev_addr << 1) | 0x01]  # 8位读地址
                    read_buf = (ctypes.c_ubyte * length)()
                    read_frame_buf = (ctypes.c_ubyte * len(read_frame))(*read_frame)
                    result2 = self.ch341.CH341StreamI2C(
                        self.device_index,
                        len(read_frame),
                        read_frame_buf,
                        length,
                        read_buf
                    )
                    return list(read_buf) if bool(result2) else None
            else:
                # 无寄存器：直接读取
                read_frame = [(dev_addr << 1) | 0x01]  # 8位读地址
                read_buf = (ctypes.c_ubyte * length)()
                read_frame_buf = (ctypes.c_ubyte * len(read_frame))(*read_frame)
                result = self.ch341.CH341StreamI2C(
                    self.device_index,
                    len(read_frame),
                    read_frame_buf,
                    length,
                    read_buf
                )
                return list(read_buf) if bool(result) else None
        except Exception as e:
            print(f"读取数据时发生错误: {e}")
            return None

    def close(self):
        """关闭设备"""
        if self.initialized and self.ch341:
            self.ch341.CH341CloseDevice(self.device_index)
            self.initialized = False


if __name__ == "__main__":
    root = tk.Tk()
    app = CH341I2CGUI(root)
    root.mainloop()